package testcase;

import org.apache.log4j.Logger;
import org.junit.Test;
import org.opengauss.jdbc.PgConnection;
import org.opengauss.jdbc.PgStatement;
import org.opengauss.quickautobalance.Cluster;
import org.opengauss.quickautobalance.ConnectionManager;
import org.opengauss.quickautobalance.DataNode;
import org.opengauss.quickautobalance.ReflectUtil;
import org.opengauss.util.GT;
import org.opengauss.util.HostSpec;
import org.opengauss.util.PSQLException;
import util.ConnectionParams;
import util.PgStatActivityUtil;

import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;


/**
 *
 */
public class QuickAutoBalanceTest {
    
    private Logger LOGGER = Logger.getLogger(QuickAutoBalanceTest.class);
    
    // 手动创建300个连接，然后重启一个节点，触发快速负载均衡，其他两个节点上的连接被逐步关掉
    @Test
    public void quickAutoBalanceAllIdleTest () throws InterruptedException, PSQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams();
        System.out.println(url);
        // 统计各节点上的连接数
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
        }
        LOGGER.info("关闭节点前:" + connectionSize);
        // 在此期间关闭数据库
        // 等待20秒后重启数据库
        while (true) {
            List<Integer> clusterConnectionSize = PgStatActivityUtil.getClusterConnectionSize(ConnectionParams.getHosts(), ConnectionParams.user,
                ConnectionParams.password, ConnectionParams.clientHost, ConnectionParams.driver, ConnectionParams.database);
            LOGGER.info(GT.tr("剩余连接数: host={0}, connectionSize={1}", ConnectionParams.getHosts().toString(),
                clusterConnectionSize.toString()));
            Thread.sleep(10*1000);
        }
    }
    
    // 创建300个连接，其中90个处于忙碌状态，重启节点，minReservedConPerDatanode=30，预期保留51个连接。
    @Test
    public void quickAutoBalanceAllIdleWithParamsTest () throws InterruptedException, PSQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams() + "&minReservedConPerCluster=20&minReservedConPerDatanode=30";
        System.out.println(url);
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        int query = 90;
        List<PgConnection> inQueryConnections = new ArrayList<>();
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);

            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
            if (i < query) {
                inQueryConnections.add(pgConnection);
            }
        }
        LOGGER.info("关闭节点前:" + connectionSize.toString());
        // 在此期间关闭数据库
        // 等待20秒后重启数据库
        while (true) {
            for (Iterator<PgConnection> iterator = inQueryConnections.iterator(); iterator.hasNext();) {
                PgConnection pgConnection = iterator.next();
                try {
                    if (!pgConnection.isValid(1)) {
                         iterator.remove();
                    }
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }
            List<Integer> clusterConnectionSize = PgStatActivityUtil.getClusterConnectionSize(ConnectionParams.getHosts(), ConnectionParams.user,
                ConnectionParams.password, ConnectionParams.clientHost, ConnectionParams.driver, ConnectionParams.database);
            LOGGER.info(GT.tr("剩余连接数: host={0}, connectionSize={1}", ConnectionParams.getHosts().toString(),
                clusterConnectionSize.toString()));
            Thread.sleep(10*1000);
        }
    }
    
    // 建300个连接，其中100个连接定时查询重启节点，minReservedConPerDatanode=30，预期保留54个连接。
    @Test
    public void quickAutoBalanceSomeIdleWithTest () throws InterruptedException, PSQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams() + "&minReservedConPerCluster=20&minReservedConPerDatanode=30";
        System.out.println(url);
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        int idle = 100;
        List<PgConnection> inQueryConnections = new ArrayList<>();
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            if (i < idle) {
                inQueryConnections.add(pgConnection);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
        }
        LOGGER.info("关闭节点前:" + connectionSize.toString());
        // 在此期间关闭数据库
        // 等待20秒后重启数据库
        while (true) {
            for (Iterator<PgConnection> iterator = inQueryConnections.iterator(); iterator.hasNext();) {
                PgConnection pgConnection = iterator.next();
                try {
                    if (pgConnection.isValid(1)) {
                        Statement statement = pgConnection.createStatement();
                        statement.execute("select 1");
                    } else {
                        iterator.remove();
                    }
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }
            List<Integer> clusterConnectionSize = PgStatActivityUtil.getClusterConnectionSize(ConnectionParams.getHosts(), ConnectionParams.user,
                ConnectionParams.password, ConnectionParams.clientHost, ConnectionParams.driver, ConnectionParams.database);
            LOGGER.info(GT.tr("剩余连接数: host={0}, connectionSize={1}", ConnectionParams.getHosts().toString(),
                clusterConnectionSize.toString()));
            Thread.sleep(10*1000);
        }
    }
    
    @Test
    public void isValidTest() throws SQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams();
        System.out.println(url);
        List<PgConnection> connections = new ArrayList<>();
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
        pgConnection.isValid(1);
        PgStatement pgStatement = (PgStatement) pgConnection.createStatement();
        pgStatement.execute("select 1");
    }
    
}
